package com.flame.system.tools.controller;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.flame.system.tools.entity.Server;
import com.flame.system.tools.entity.SshMessage;
import com.flame.system.tools.mapper.ServerMapper;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.Session;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * ServerSshController
 *
 * @author wuxintong
 * @since 2022/11/26
 */
@Service("server-ssh-controller")
public class ServerSshController extends TextWebSocketHandler {

    @Autowired
    private ObjectMapper objectMapper;

    @Autowired
    private ServerMapper serverMapper;

    private Map<String, WebSocketSession> sessionMap = new HashMap<>();

    private Map<String, Session> jSchMap = new HashMap<>();

    private Map<String, Channel> channelMap = new HashMap<>();

    private ExecutorService executorService = Executors.newCachedThreadPool();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {

    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        SshMessage sshMessage = objectMapper.readValue(message.getPayload(), SshMessage.class);
        switch (sshMessage.getCommandType()) {
            case "CONNECT":
                connect(session, sshMessage);
                break;
            case "SHELL":
                shell(session, sshMessage);
                break;
            case "CLOSE":
                close(session);
                break;
            default:
                break;
        }
    }

    private void connect(WebSocketSession webSocketSession, SshMessage sshMessage) {
        try {
            Server server = serverMapper.selectById(sshMessage.getServerName());
            JSch jSch = new JSch();
            Session jSchSession =
                    jSch.getSession(server.getServerUsername(), server.getServerHost(), server.getServerPort());
            jSchSession.setConfig("StrictHostKeyChecking", "no");
            jSchSession.setPassword(server.getServerPassword());
            jSchSession.connect(30000);
            Channel channel = jSchSession.openChannel("shell");
            channel.connect(3000);

            executorService.execute(() -> {
                try (InputStream inputStream = channel.getInputStream()) {
                    byte[] buf = new byte[1024];
                    int len = 0;
                    while ((len = inputStream.read(buf)) != -1) {
                        String message = new String(buf, 0, len, StandardCharsets.UTF_8);
                        webSocketSession.sendMessage(new TextMessage(message));
                    }
                } catch (Exception e) {
                    close(webSocketSession);
                }
            });
            webSocketSession.sendMessage(new TextMessage("连接成功\r\n"));
            webSocketSession.getAttributes().put("account", sshMessage.getAccount());
            sessionMap.put(sshMessage.getAccount(), webSocketSession);
            jSchMap.put(sshMessage.getAccount(), jSchSession);
            channelMap.put(sshMessage.getAccount(), channel);
        } catch (Exception e) {
            close(webSocketSession);
        }
    }

    private void shell(WebSocketSession session, SshMessage sshMessage) throws IOException {
        Channel channel = channelMap.get(sshMessage.getAccount());
        if (channel != null) {
            OutputStream outputStream = channel.getOutputStream();
            outputStream.write(sshMessage.getCommand().getBytes(StandardCharsets.UTF_8));
            outputStream.flush();
        }
    }

    private void close(WebSocketSession session) {
        if (session != null) {
            String account = (String) session.getAttributes().get("account");
            Channel channel = channelMap.get(account);
            if (channel != null) {
                channel.disconnect();
                channelMap.remove(account);
            }

            Session jschSession = jSchMap.get(account);
            if (jschSession != null) {
                jschSession.disconnect();
                jSchMap.remove(account);
            }

            WebSocketSession webSocketSession = sessionMap.get(account);
            if (webSocketSession != null) {
                try {
                    webSocketSession.close();
                    sessionMap.remove(account);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        close(session);
    }
}
